#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Aug 19 14:07:39 2021

@author: ljl
"""


import tensorflow as tf
import cv2 as cv
import numpy as np
import sys
import os
import glob
import copy
from skimage import io, transform
import matplotlib.pyplot as plt

#import stat
#import shutil

row = 28
col = 28
c = 3
img_type = '.png'


def loadImgNameAB(rowDataPath):
    sampleA = []
    sampleB = []
    for file in os.listdir(rowDataPath):
        idxA = file.find('a')
        idxB = file.find('b')
        if ( idxA != -1):
            #a = file[idxA]
            sampleA.append(rowDataPath + "/" + file)
        if (idxB != -1):
            #b = file[idxB]
            sampleB.append(rowDataPath + "/" + file)
        
    return sampleA, sampleB

def loadImgName(inPath):
    imgPath = []
    name = []
    for file in os.listdir(inPath):
        imgPath.append(inPath + "/" + file)
        name.append(file)      #xxx.png
    return imgPath, name


#label: T or F
def mergeImg(path, label, savePath):
    imgPathA, imgPathB = loadImgNameAB(path)
    imgPathA.sort()
    imgPathB.sort()
    if(len(imgPathA) != len(imgPathB)):
        print("a, b len not equal")
        sys.exit()
    print("sample number: %d", len(imgPathA))
    
    count = 0
    for nameA in imgPathA:
        idx = nameA.find('a.png')
        tempB = nameA[0:idx]
        tempB = tempB + 'b.png'
        nameB = 'null'

        if os.path.exists(tempB):
            nameB = tempB
        else:
            print("path not exist", tempB)
            sys.exit()
        if nameB is 'null':
            print("%s is not find in nameB", nameA)
            sys.exit()
            
        if os.path.exists(nameA) and os.path.exists(nameB):
            imgA = cv.imread(nameA, 0)
            imgB = cv.imread(nameB, 0)

            print(imgA.shape)
            print(imgB.shape)
            zeros = np.zeros(imgA.shape[0:2], dtype="uint8")
            mergeImg = cv.merge( [imgA, imgB, zeros]  )
            
            cv.imwrite(savePath + "/" + label + str(count) + ".jpg", mergeImg)
            count += 1
                
   
#[n, row, col, c]
def generateDataSet(mergeImgPath, row, col):
    
    imgPathList, nameList = loadImgName(mergeImgPath)
    imgPathList.sort()
    nameList.sort()
    imgList = []
    lableList = []
    
    #label = np.zeros((1,2))
    size = len(imgPathList)
    for idx in range(size):
        imgTemp = cv.imread(imgPathList[idx])
        imgTemp = transform.resize(imgTemp, (row, col, imgTemp.shape[2]), mode='constant')
        
        imgList.append(imgTemp[:,:,0:2])
     
        if nameList[idx][0] is 'T':
            label = [0, 1]
            lableList.append(label)
        else:
            #print(nameList[idx])
            label = [1, 0]
            lableList.append(label)
            
    return np.asarray(imgList, np.float32), np.asarray(lableList, np.int32)

#0:diff  1:sim
#[1, 0],  [0,1]
def one_hot(labels):
    labels_hot = []
    for label in labels:
        if label == 0:
            labels_hot.append([1,0])
        else:
            labels_hot.append([0,1])
    return np.asarray(labels_hot, np.int32)

def _show_time(cost_time):
    #start_time = time.time()
    hour =   ( cost_time // 60 ) // 60
    minute = ( (cost_time) // 60 ) % 60
    second = (  cost_time) % 60
    
    print('Running time:%f Second' % cost_time)  # 输出运行时间
    print('Running time: {:.0f}h {:.0f}m {:.0f}s '.format(hour, minute, second))
    return hour, minute, second


def slipImgChannel(imgs):
    if len(imgs.shape) != 4:
        print("imgs dims error: ", imgs.shape)
        sys.exit(0)
    a,b,c,d = imgs.shape
    img0_list = []
    img1_list = []
 
    for batch_img in range(a):
        temp1 = np.ones((b, c, 1)) 
        temp2 = np.ones((b, c, 1))
        
        temp1[:,:,0] = imgs[batch_img][:,:,0]
        img0_list.append(temp1)
        
        temp2[:,:,0] = imgs[batch_img][:,:,1]
        img1_list.append(temp2)
        
        
    return np.asarray(img0_list, np.float32), np.asarray(img1_list, np.float32)

def read_img(path):
    # 0  4   2   3   1
    cate = [x for x in os.listdir(path) if os.path.isdir(path + x)]      #0   1   2   3   4   5
    cate.sort()
    
    print('------------', cate)
    imgs = []
    labels = []
    np_label = np.zeros((1, len(cate)), dtype=np.int)
    org_label = list(np_label[0])
    count = 0
    for idx, folder in enumerate(cate):
        print ('\n dir is ', folder)
        for im in glob.glob(path + folder + '/*' + img_type):
            count += 1
            print('\r'+'reading the image %d: %s' % (count,im), end='',flush=True)
            make_label = copy.deepcopy(org_label)
            make_label[idx] = 1
            # 0:dian_gui58-842   1:filter_things154-852   2:铁spots132-856   3：ping_tai857   4：rail371-835  5:ground41-818  6:line179-806  
            #7/11:foreign1-718  8/12:foreign2-750

            labels.append(make_label)

            img = cv.imread(im)
            img = transform.resize(img, (row, col, c), mode='constant')
            imgs.append(img)
    return np.asarray(imgs, np.float32), np.asarray(labels, np.int32)

def showLoss(cossList, cossList2):
# 绘制损失曲线
    #plt.clf()
    plt.plot(cossList, 'g-', label='train_loss')
    plt.plot(cossList2, 'b-', label='valid_loss')
    
    plt.title('model loss')
    plt.xlabel('epoch')
    plt.ylabel('loss value')
    plt.legend(['train', 'valid'], loc='upper left')
    
    #plt.clf()
    #plt.plot(cost, 'g-')
    plt.pause(0.01)

def showAccuracy(accList, accList2):
# 绘制损失曲线
    plt.clf()
    plt.plot(accList, 'g-', label='train_accuracy')
    plt.plot(accList2, 'b-', label='valid_accuracy')
    
    plt.title('model accuracy')
    plt.xlabel('epoch')
    plt.ylabel('accuracy value')
    plt.legend(['train', 'valid'], loc='upper left')
    
    #plt.clf()
    #plt.plot(cost, 'g-')
    plt.pause(0.01)
    
def smoothVal(last, curr, weight = 0.7):

    if last == -1: last = curr
    smoothed = last * weight + (1 - weight)*curr
    return smoothed
    

def shuffleDatas(imgs, labels):
# 打乱顺序
    num_example = imgs.shape[0]
    arr = np.arange(num_example)  # arr = [0,1,2,........]
    np.random.shuffle(arr)
    return imgs[arr] , labels[arr]  # 打乱顺序
    

if __name__ == '__main__':
    
#    mergeImg("/home/ljl/Desktop/newImagesData/1", 'T', "/home/ljl/Desktop/newImagesData/merge")
#    mergeImg("/home/ljl/Desktop/newImagesData/0", 'F', "/home/ljl/Desktop/newImagesData/merge")
    
    imgs, labels = generateDataSet("/home/ljl/Desktop/newImagesData/merge", row,col)
    print(imgs.shape)
    print(labels.shape)

    sys.exit()
    batch_left, batch_right = slipImgChannel(imgs)
    print(batch_left.shape)

    start = 0 
    for idx in range(2):
        print (idx)
        print(batch_left.shape)
        cv.imshow("aaa", batch_left[idx + start])
        cv.imshow("bbb", batch_right[idx + start])
        cv.waitKey(0)
    cv.destroyAllWindows()
    sys.exit()

    for idx in range(20):
        print (idx)
        cv.imshow("aaa", imgs[idx + start][:,:,0])
        cv.imshow("bbb", imgs[idx + start][:,:,1])
        cv.waitKey(0)
    cv.destroyAllWindows()

    